//
// Created by 郝羽 on 2018/2/5.
//

#ifndef RTC_XSWITCH_MIXER_OUTPUT_HPP
#define RTC_XSWITCH_MIXER_OUTPUT_HPP

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
};

class Output {
public:
    int Open(const std::string& file_name, const std::string& format_name, const std::string& codec_name) {
        AVCodec* codec = avcodec_find_encoder_by_name(codec_name.c_str());
        if (!codec) {
            std::clog << "could not find " << codec_name << " encoder" << std::endl;
            return -1;
        }

        encoder_ = avcodec_alloc_context3(codec);
        if (!encoder_) {
            std::clog << "could not allocate encoder" << std::endl;
            return -1;
        }

        encoder_->codec_id = codec->id;
        encoder_->codec_type = codec->type;
        encoder_->sample_rate = 16000;
        encoder_->sample_fmt = AV_SAMPLE_FMT_S16;
        encoder_->channel_layout = AV_CH_LAYOUT_MONO;
        encoder_->channels = 1;

        int ret = avcodec_open2(encoder_, codec, nullptr);
        if (ret < 0) {
            std::clog << "could not open " << codec_name << " encoder: " << av_err2str(ret);
            return -1;
        }

        ret = avformat_alloc_output_context2(&fmt_, nullptr, format_name.c_str(), nullptr);
        if (ret < 0) {
            std::clog << "could not allocate output format " << av_err2str(ret) << std::endl;
            avcodec_free_context(&encoder_);
            return -1;
        }

        AVStream* stream = avformat_new_stream(fmt_, codec);
        stream->codecpar->format = AV_SAMPLE_FMT_S16;
        stream->codecpar->channels = 1;
        stream->codecpar->channel_layout = AV_CH_LAYOUT_MONO;
        stream->codecpar->sample_rate = 16000;
        stream->codecpar->codec_id = codec->id;
        stream->codecpar->codec_type = codec->type;

        ret = avio_open(&fmt_->pb, file_name.c_str(), AVIO_FLAG_WRITE);
        if (ret < 0) {
            std::clog << "could not open " << file_name << av_err2str(ret) << std::endl;
            avcodec_free_context(&encoder_);
            avformat_free_context(fmt_);
            return -1;
        }

        ret = avformat_write_header(fmt_, nullptr);
        if (ret < 0) {
            std::clog << "could not write header" << av_err2str(ret) << std::endl;
            avcodec_free_context(&encoder_);
            avformat_free_context(fmt_);
            avio_close(fmt_->pb);
            return -1;
        }

        return 0;
    }

    int Write(AVFrame* frame) {
        int ret = avcodec_send_frame(encoder_, frame);
        if (ret < 0) {
            std::clog << "output error: " << av_err2str(ret) << std::endl;
            return -1;
        }

        AVPacket packet;
        av_init_packet(&packet);
        while (avcodec_receive_packet(encoder_, &packet) == 0) {
            ret = av_interleaved_write_frame(fmt_, &packet);
            if (ret < 0) {
                std::clog << "output error: " << av_err2str(ret) << std::endl;
                return -1;
            }
        }
    }

    int Close() {
        av_write_trailer(fmt_);

        if (fmt_) {
            if (fmt_->pb) {
                avio_close(fmt_->pb);
                fmt_->pb = nullptr;
            }
            avformat_free_context(fmt_);
            fmt_ = nullptr;
        }
        if (encoder_) {
            avcodec_free_context(&encoder_);
            encoder_ = nullptr;
        }
    }

private:
    AVFormatContext* fmt_ = nullptr;
    AVCodecContext* encoder_ = nullptr;
};

#endif //RTC_XSWITCH_MIXER_OUTPUT_HPP
